﻿# coding=utf-8
# 2023/4/27: 首个记录，基于旧2020/11/6版本
# 2023/5/5: 枚举基于Enum
# 2023/9/12: StopSign改为TrafficSign，并增加TrafficSignType
# 2023/10/11: 补充部分英文注释
# 2024/5/20: 新增TrafficLight.stop_pos_x/y
# 2025/3/26: 实现样本插值
# 2025/7/18: 修正significant赋值
# 2025/7/21: 使用类型注解。优化to_general_sample

import bi_common as bi
from enum import Enum
from typing import List, Any


class TrafficElementType(Enum):
    INVALID = 0  # 无效
    SPEED_LIMIT = 1  # 限速标志
    TRAFFIC_SIGN = 2  # 交通标志
    TRAFFIC_LIGHT = 3  # 交通信号灯
    STOP_LINE = 4  # 停止线
    SPEED_BUMP = 5  # 减速带
    ARROW_MARK = 6  # 路面箭头
    PROHIBITION_ZONE = 7  # 禁停区
    PARKING_SLOT = 8  # 停车位
    ZEBRA_CROSSING = 9  # 斑马线


class TrafficSignType(Enum):
    STOP = 0,  # 停止标志

    NO_ENTRY = 1,  # 禁止驶入
    NO_DRIVING = 2,  # 禁止通行
    No_Parking = 3,  # 禁止停车
    No_LONG_TIME_PARKING = 4,  # 禁止长时间停车
    NO_VEHICLE = 5,  # 禁止机动车驶入
    NO_TRUCK = 6,  # 禁止货车驶入
    NO_BUS = 7,  # 禁止客车驶入
    NO_PEDESTRIAN = 8,  # 禁止行人进入
    NO_LEFT_TURN = 9,  # 禁止左转
    NO_RIGHT_TURN = 10,  # 禁止右转
    NO_ANY_TURN = 11,  # 禁止转弯
    NO_FORWARD = 12,  # 禁止直行
    NO_FORWARD_OR_LEFT = 13,  # 禁止直行或左转
    NO_FORWARD_OR_RIGHT = 14,  # 禁止直行或右转
    NO_U_TURN = 15,  # 禁止调头
    NO_HORN = 16,  # 禁止鸣笛
    NO_OVERTAKE = 17,  # 禁止超车
    NO_OVERTAKE_ENDER = 18,  # 禁止超车区间结束
    NO_DANGEROUS_GOODS = 19,  # 禁止危险品车辆驶入

    LOW_SPEED_LIMIT = 50,  # 低速限速
    SPEED_LIMIT_ENDER = 51,  # 限速区间结束
    WEIGHT_LIMIT_WHEEL = 52,  # 轴重限重
    WEIGHT_LIMIT_TOTAL = 53,  # 全车限重
    HEIGHT_LIMIT = 54,  # 高度限制
    WIDTH_LIMIT = 55,  # 宽度限制
    YIELD = 56,  # 让行标志
    STOP_AND_CHECK = 57,  # 停车检查

    GENERAL_SIGN = 100,  # 指示类标志
    GENERAL_NO_SPECIAL_VEHICLE = 101,  # 禁止特种车辆驶入
    GENERAL_CONSTRUCTION = 102,  # 施工标志
    GENERAL_WARNING = 103,  # 警告牌
    GENERAL_GREEN_BACKGROUND = 104,  # 绿底指示牌
    GENERAL_BLUE_BACKGROUND = 105,  # 蓝底指示牌
    GENERAL_RED_BACKGROUND = 106,  # 红底指示牌
    GENERAL_WHITE_BACKGROUND = 107,  # 白底指示牌
    GENERAL_YELLOW_ROAD_ID = 108,  # 黄底道路ID牌
    GENERAL_RED_ROAD_ID = 109,  # 红底道路ID牌


class TrafficLightState(Enum):
    INVALID = 0  # 无效
    STOP = 1  # 红灯
    CONTINUE = 2  # 黄灯
    CAUTION = 3  # 黄灯闪烁
    ALLOW = 4  # 绿灯
    ALLOW_FLASH = 5  # 绿灯闪烁
    STOP_FLASH = 6  # 红灯闪烁


class ParkingSlotOccupancy(Enum):
    UNKNOWN = 0  # 未知
    NOT_OCCUPIED = 1  # 未被占用
    OCCUPIED = 2  # 被占用


class ArrowCombination:
    def __init__(self):
        self.valid: bool = False  # False while direction is not specified / 是否有效，若不分方向(如实心圆交通灯信号)则为False
        self.go_straight: bool = False  # 直行
        self.turn_left: bool = False  # 左转
        self.turn_right: bool = False  # 右转
        self.u_turn: bool = False  # 调头


# Base class of traffic element / 交通元素基类
class TrafficElement:
    def __init__(self):
        self.id: int = 0  # 0 as invalid / 交通元素ID，0表示无效
        self.raw_id: int | None = None  # 交通元素的原始ID


class SpeedLimit(TrafficElement):
    def __init__(self):
        TrafficElement.__init__(self)
        self.limit_speed: int = 0  # [KPH] 限速
        self.pos_x: float | None = None  # [m] 纵向位置
        self.pos_y: float | None = None  # [m] 横向位置
        self.pos_z: float | None = None  # [m] 垂直位置


class TrafficSign(TrafficElement):
    def __init__(self):
        TrafficElement.__init__(self)
        self.sign_type: TrafficSignType = TrafficSignType.STOP  # 交通标志类型
        self.pos_x: float | None = None  # [m] 纵向位置
        self.pos_y: float | None = None  # [m] 横向位置
        self.pos_z: float | None = None  # [m] 垂直位置


class TrafficLight(TrafficElement):
    def __init__(self):
        TrafficElement.__init__(self)
        self.state: TrafficLightState = TrafficLightState.INVALID  # 信号灯状态
        self.arrows: ArrowCombination = ArrowCombination()  # False while direction is not specified / 信号灯的箭头组合，若不分方向(实心圆信号)则valid为False
        self.pos_x: float | None = None  # [m] 纵向位置
        self.pos_y: float | None = None  # [m] 横向位置
        self.pos_z: float | None = None  # [m] 垂直位置
        self.stop_pos_x: float | None = None  # [m] 纵向停止位置
        self.stop_pos_y: float | None = None  # [m] 横向停止位置


class StopLine(TrafficElement):
    def __init__(self):
        TrafficElement.__init__(self)
        self.distance: float = 0.0  # [m] 停止线与x轴相交位置
        self.heading: float | None = None  # [°] 停止线法线朝向，逆时针为正方向


class SpeedBump(TrafficElement):
    def __init__(self):
        TrafficElement.__init__(self)
        self.distance: float = 0.0  # [m] 减速带与x轴相交位置
        self.heading: float | None = None  # [°] 减速带法线朝向，逆时针为正方向


class ArrowMark(TrafficElement):
    def __init__(self):
        TrafficElement.__init__(self)
        self.arrows: ArrowCombination = ArrowCombination()  # 路面标志的箭头组合
        self.pos_x: float | None = None  # [m] 纵向位置
        self.pos_y: float | None = None  # [m] 横向位置


class ProhibitionZone(TrafficElement):
    def __init__(self):
        TrafficElement.__init__(self)
        self.center_x: float = 0.0  # [m] 禁停区中央的纵向位置
        self.center_y: float = 0.0  # [m] 禁停区中央的横向位置
        self.width: float = 0.0  # [m] 禁停区的宽度
        self.length: float = 0.0  # [m] 禁停区的长度
        self.heading: float | None = None  # [°] 禁停区的朝向(沿长度方向)，逆时针为正方向


class ParkingSlot(TrafficElement):
    def __init__(self):
        TrafficElement.__init__(self)
        self.center_x: float = 0.0  # [m] 停车位中央的纵向位置
        self.center_y: float = 0.0  # [m] 停车位中央的横向位置
        self.width: float = 0.0  # [m] 停车位的宽度
        self.length: float = 0.0  # [m] 停车位的长度
        self.heading: float | None = None  # [°] 停车位的朝向(沿长度方向)，逆时针为正方向
        self.occupancy: ParkingSlotOccupancy = ParkingSlotOccupancy.UNKNOWN  # 停车位是否被占用


class ZebraCrossing(TrafficElement):
    def __init__(self):
        TrafficElement.__init__(self)
        self.near_distance: float = 0.0  # [m] 斑马线近端与x轴相交位置
        self.far_distance: float | None = None  # [m] 斑马线远端与x轴相交位置
        self.heading: float | None = None  # [°] 斑马线的朝向，逆时针为正方向


class TrafficSample:
    def __init__(self):
        self.time: float = 0.0  # [s] Time offset in session / Session内的相对时间
        self.elements: List[TrafficElement] = []  # List of traffic elements / 交通元素列表

    # Convert to general sample for output / 转通用样本，用于样本输出
    def to_general_sample(self, channel: int) -> Any:
        output = bi.agency.create_general_sample()
        output.protocol = "traffic-sample-v2@" + str(channel)
        output.time = self.time
        types = []
        offsets = []
        sizes = []
        datas = []
        total_data_size = 0
        for elem in self.elements:
            elem_data = []
            if isinstance(elem, SpeedLimit):
                types.append(TrafficElementType.SPEED_LIMIT)
                elem_data.append(elem.limit_speed)
                elem_data.append(elem.pos_x)
                elem_data.append(elem.pos_y)
                elem_data.append(elem.pos_z)
                elem_data.append(None)  # Reserved
                elem_data.append(None)  # Reserved
            elif isinstance(elem, TrafficSign):
                types.append(TrafficElementType.TRAFFIC_SIGN)
                elem_data.append(elem.pos_x)
                elem_data.append(elem.pos_y)
                elem_data.append(elem.pos_z)
                elem_data.append(elem.sign_type.value)
                elem_data.append(None)  # Reserved
            elif isinstance(elem, TrafficLight):
                types.append(TrafficElementType.TRAFFIC_LIGHT)
                elem_data.append(elem.state.value)
                elem_data.append(elem.pos_x)
                elem_data.append(elem.pos_y)
                elem_data.append(elem.pos_z)
                if elem.arrows.valid:
                    elem_data.append(elem.arrows.go_straight)
                    elem_data.append(elem.arrows.turn_left)
                    elem_data.append(elem.arrows.turn_right)
                    elem_data.append(elem.arrows.u_turn)
                else:
                    elem_data.append(None)
                    elem_data.append(None)
                    elem_data.append(None)
                    elem_data.append(None)
                elem_data.append(elem.stop_pos_x)
                elem_data.append(elem.stop_pos_y)
            elif isinstance(elem, StopLine):
                types.append(TrafficElementType.STOP_LINE)
                elem_data.append(elem.distance)
                elem_data.append(elem.heading)
                elem_data.append(None)  # Reserved
                elem_data.append(None)  # Reserved
            elif isinstance(elem, SpeedBump):
                types.append(TrafficElementType.SPEED_BUMP)
                elem_data.append(elem.distance)
                elem_data.append(elem.heading)
                elem_data.append(None)  # Reserved
                elem_data.append(None)  # Reserved
            elif isinstance(elem, ArrowMark):
                types.append(TrafficElementType.ARROW_MARK)
                elem_data.append(elem.arrows.go_straight)
                elem_data.append(elem.arrows.turn_left)
                elem_data.append(elem.arrows.turn_right)
                elem_data.append(elem.arrows.u_turn)
                elem_data.append(elem.pos_x)
                elem_data.append(elem.pos_y)
                elem_data.append(None)  # Reserved
                elem_data.append(None)  # Reserved
            elif isinstance(elem, ProhibitionZone):
                types.append(TrafficElementType.PROHIBITION_ZONE)
                elem_data.append(elem.center_x)
                elem_data.append(elem.center_y)
                elem_data.append(elem.width)
                elem_data.append(elem.length)
                elem_data.append(elem.heading)
                elem_data.append(None)  # Reserved
                elem_data.append(None)  # Reserved
            elif isinstance(elem, ParkingSlot):
                types.append(TrafficElementType.PARKING_SLOT)
                elem_data.append(elem.center_x)
                elem_data.append(elem.center_y)
                elem_data.append(elem.width)
                elem_data.append(elem.length)
                elem_data.append(elem.heading)
                elem_data.append(elem.occupancy.value)
                elem_data.append(None)  # Reserved
            elif isinstance(elem, ZebraCrossing):
                types.append(TrafficElementType.ZEBRA_CROSSING)
                elem_data.append(elem.near_distance)
                elem_data.append(elem.far_distance)
                elem_data.append(elem.heading)
                elem_data.append(None)  # Reserved
                elem_data.append(None)  # Reserved
            if len(elem_data) > 0:
                offsets.append(total_data_size)
                sizes.append(len(elem_data))
                datas.append(elem_data)
                total_data_size += len(elem_data)
        
        total_length = 1 + len(self.elements) * 6 + total_data_size
        output.significant = total_length
        
        values: List[float | str | None] = [None] * total_length
        values[0] = len(self.elements)
        
        for i in range(0, len(self.elements)):
            base_index = 1 + i * 6
            values[base_index] = offsets[i]
            values[base_index + 1] = sizes[i]
            values[base_index + 2] = types[i].value
            values[base_index + 3] = self.elements[i].id
            values[base_index + 4] = self.elements[i].raw_id
            # 索引5为保留字段，保持None值
        
        data_index = 1 + len(self.elements) * 6
        for elem_data in datas:
            for j, data in enumerate(elem_data):
                values[data_index + j] = data
            data_index += len(elem_data)
        
        output.values = values
        return output


def _conv_traffic_sample_v1(gs: Any) -> TrafficSample | None:
    values_count = len(gs.values)
    if values_count < 1:
        return None
    elem_count = int(gs.values[0]) if gs.values[0] is not None else 0
    if values_count < 1 + elem_count * 3:
        return None
    output = TrafficSample()
    output.time = gs.time
    for i in range(0, elem_count):
        if gs.values[1 + 3 * i] is None or gs.values[2 + 3 * i] is None or gs.values[3 + 3 * i] is None:
            continue
        type = TrafficElementType(int(gs.values[1 + 3 * i]))
        eo = int(gs.values[2 + 3 * i]) + 1 + elem_count * 3
        es = int(gs.values[3 + 3 * i])
        if values_count < eo + es:
            return None
        elem = None
        if type == TrafficElementType.SPEED_LIMIT:
            if gs.values[eo] is not None:
                elem = SpeedLimit()
                elem.limit_speed = int(gs.values[eo])
                elem.pos_x = float(gs.values[eo + 1]) if gs.values[eo + 1] is not None else None
                elem.pos_y = float(gs.values[eo + 2]) if gs.values[eo + 2] is not None else None
        elif type == TrafficElementType.TRAFFIC_SIGN:
            elem = TrafficSign()
            elem.pos_x = float(gs.values[eo]) if gs.values[eo] is not None else None
            elem.pos_y = float(gs.values[eo + 1]) if gs.values[eo + 1] is not None else None
        elif type == TrafficElementType.TRAFFIC_LIGHT:
            if gs.values[eo] is not None:
                elem = TrafficLight()
                elem.state = TrafficLightState(int(gs.values[eo]))
                elem.pos_x = float(gs.values[eo + 1]) if gs.values[eo + 1] is not None else None
                elem.pos_y = float(gs.values[eo + 2]) if gs.values[eo + 2] is not None else None
                if gs.values[eo + 3] is not None and gs.values[eo + 4] is not None and gs.values[eo + 5] is not None and gs.values[eo + 6] is not None:
                    elem.arrows.valid = True
                    elem.arrows.go_straight = gs.values[eo + 3] == 1
                    elem.arrows.turn_left = gs.values[eo + 4] == 1
                    elem.arrows.turn_right = gs.values[eo + 5] == 1
                    elem.arrows.u_turn = gs.values[eo + 6] == 1
        elif type == TrafficElementType.STOP_LINE:
            if gs.values[eo] is not None:
                elem = StopLine()
                elem.distance = float(gs.values[eo])
                elem.heading = float(gs.values[eo + 1]) if gs.values[eo + 1] is not None else None
        elif type == TrafficElementType.SPEED_BUMP:
            if gs.values[eo] is not None:
                elem = SpeedBump()
                elem.distance = float(gs.values[eo])
                elem.heading = float(gs.values[eo + 1]) if gs.values[eo + 1] is not None else None
        elif type == TrafficElementType.ARROW_MARK:
            if gs.values[eo] is not None and gs.values[eo + 1] is not None and gs.values[eo + 2] is not None and gs.values[eo + 3] is not None:
                elem = ArrowMark()
                elem.arrows.valid = True
                elem.arrows.go_straight = gs.values[eo] == 1
                elem.arrows.turn_left = gs.values[eo + 1] == 1
                elem.arrows.turn_right = gs.values[eo + 2] == 1
                elem.arrows.u_turn = gs.values[eo + 3] == 1
                elem.pos_x = float(gs.values[eo + 4]) if gs.values[eo + 4] is not None else None
                elem.pos_y = float(gs.values[eo + 5]) if gs.values[eo + 5] is not None else None
        elif type == TrafficElementType.PROHIBITION_ZONE:
            if gs.values[eo] is not None and gs.values[eo + 1] is not None and gs.values[eo + 2] is not None and gs.values[eo + 3] is not None:
                elem = ProhibitionZone()
                elem.center_x = float(gs.values[eo])
                elem.center_y = float(gs.values[eo + 1])
                elem.width = float(gs.values[eo + 2])
                elem.length = float(gs.values[eo + 3])
                elem.heading = float(gs.values[eo + 4]) if gs.values[eo + 4] is not None else None
        elif type == TrafficElementType.PARKING_SLOT:
            if gs.values[eo] is not None and gs.values[eo + 1] is not None and gs.values[eo + 2] is not None and gs.values[eo + 3] is not None:
                elem = ParkingSlot()
                elem.center_x = float(gs.values[eo])
                elem.center_y = float(gs.values[eo + 1])
                elem.width = float(gs.values[eo + 2])
                elem.length = float(gs.values[eo + 3])
                elem.heading = float(gs.values[eo + 4]) if gs.values[eo + 4] is not None else None
        elif type == TrafficElementType.ZEBRA_CROSSING:
            if gs.values[eo] is not None:
                elem = ZebraCrossing()
                elem.near_distance = float(gs.values[eo])
                elem.far_distance = float(gs.values[eo + 1]) if gs.values[eo + 1] is not None else None
                elem.heading = float(gs.values[eo + 2]) if gs.values[eo + 2] is not None else None
        if elem is not None:
            output.elements.append(elem)
    return output


def _conv_traffic_sample_v2(gs: Any) -> TrafficSample | None:
    values_count = len(gs.values)
    if values_count < 1:
        return None
    elem_count = int(gs.values[0]) if gs.values[0] is not None else 0
    if values_count < 1 + elem_count * 6:
        return None
    output = TrafficSample()
    output.time = gs.time
    for i in range(0, elem_count):
        if gs.values[1 + 6 * i] is None or gs.values[2 + 6 * i] is None or gs.values[3 + 6 * i] is None or gs.values[4 + 6 * i] is None:
            continue
        eo = int(gs.values[1 + 6 * i]) + 1 + elem_count * 6
        es = int(gs.values[2 + 6 * i])
        type = TrafficElementType(int(gs.values[3 + 6 * i]))
        if values_count < eo + es:
            return None
        elem = None
        if type == TrafficElementType.SPEED_LIMIT:
            if gs.values[eo] is not None:
                elem = SpeedLimit()
                elem.limit_speed = int(gs.values[eo])
                elem.pos_x = float(gs.values[eo + 1]) if gs.values[eo + 1] is not None else None
                elem.pos_y = float(gs.values[eo + 2]) if gs.values[eo + 2] is not None else None
                elem.pos_z = float(gs.values[eo + 3]) if gs.values[eo + 3] is not None else None
        elif type == TrafficElementType.TRAFFIC_SIGN:
            elem = TrafficSign()
            elem.pos_x = float(gs.values[eo]) if gs.values[eo] is not None else None
            elem.pos_y = float(gs.values[eo + 1]) if gs.values[eo + 1] is not None else None
            elem.pos_z = float(gs.values[eo + 2]) if gs.values[eo + 2] is not None else None
            elem.sign_type = TrafficSignType(int(gs.values[eo + 3])) if gs.values[eo + 3] is not None else TrafficSignType.STOP
        elif type == TrafficElementType.TRAFFIC_LIGHT:
            if gs.values[eo] is not None:
                elem = TrafficLight()
                elem.state = TrafficLightState(int(gs.values[eo]))
                elem.pos_x = float(gs.values[eo + 1]) if gs.values[eo + 1] is not None else None
                elem.pos_y = float(gs.values[eo + 2]) if gs.values[eo + 2] is not None else None
                elem.pos_z = float(gs.values[eo + 3]) if gs.values[eo + 3] is not None else None
                if gs.values[eo + 4] is not None and gs.values[eo + 5] is not None and gs.values[eo + 6] is not None and gs.values[eo + 7] is not None:
                    elem.arrows.valid = True
                    elem.arrows.go_straight = gs.values[eo + 4] == 1
                    elem.arrows.turn_left = gs.values[eo + 5] == 1
                    elem.arrows.turn_right = gs.values[eo + 6] == 1
                    elem.arrows.u_turn = gs.values[eo + 7] == 1
                elem.stop_pos_x = float(gs.values[eo + 8]) if gs.values[eo + 8] is not None else None
                elem.stop_pos_y = float(gs.values[eo + 9]) if gs.values[eo + 9] is not None else None
        elif type == TrafficElementType.STOP_LINE:
            if gs.values[eo] is not None:
                elem = StopLine()
                elem.distance = float(gs.values[eo])
                elem.heading = float(gs.values[eo + 1]) if gs.values[eo + 1] is not None else None
        elif type == TrafficElementType.SPEED_BUMP:
            if gs.values[eo] is not None:
                elem = SpeedBump()
                elem.distance = float(gs.values[eo])
                elem.heading = float(gs.values[eo + 1]) if gs.values[eo + 1] is not None else None
        elif type == TrafficElementType.ARROW_MARK:
            if gs.values[eo] is not None and gs.values[eo + 1] is not None and gs.values[eo + 2] is not None and gs.values[eo + 3] is not None:
                elem = ArrowMark()
                elem.arrows.valid = True
                elem.arrows.go_straight = gs.values[eo] == 1
                elem.arrows.turn_left = gs.values[eo + 1] == 1
                elem.arrows.turn_right = gs.values[eo + 2] == 1
                elem.arrows.u_turn = gs.values[eo + 3] == 1
                elem.pos_x = float(gs.values[eo + 4]) if gs.values[eo + 4] is not None else None
                elem.pos_y = float(gs.values[eo + 5]) if gs.values[eo + 5] is not None else None
        elif type == TrafficElementType.PROHIBITION_ZONE:
            if gs.values[eo] is not None and gs.values[eo + 1] is not None and gs.values[eo + 2] is not None and gs.values[eo + 3] is not None:
                elem = ProhibitionZone()
                elem.center_x = float(gs.values[eo])
                elem.center_y = float(gs.values[eo + 1])
                elem.width = float(gs.values[eo + 2])
                elem.length = float(gs.values[eo + 3])
                elem.heading = float(gs.values[eo + 4]) if gs.values[eo + 4] is not None else None
        elif type == TrafficElementType.PARKING_SLOT:
            if gs.values[eo] is not None and gs.values[eo + 1] is not None and gs.values[eo + 2] is not None and gs.values[eo + 3] is not None:
                elem = ParkingSlot()
                elem.center_x = float(gs.values[eo])
                elem.center_y = float(gs.values[eo + 1])
                elem.width = float(gs.values[eo + 2])
                elem.length = float(gs.values[eo + 3])
                elem.heading = float(gs.values[eo + 4]) if gs.values[eo + 4] is not None else None
                elem.occupancy = ParkingSlotOccupancy(int(gs.values[eo + 5])) if gs.values[eo + 5] is not None else ParkingSlotOccupancy.UNKNOWN
        elif type == TrafficElementType.ZEBRA_CROSSING:
            if gs.values[eo] is not None:
                elem = ZebraCrossing()
                elem.near_distance = float(gs.values[eo])
                elem.far_distance = float(gs.values[eo + 1]) if gs.values[eo + 1] is not None else None
                elem.heading = float(gs.values[eo + 2]) if gs.values[eo + 2] is not None else None
        if elem is not None:
            elem.id = int(gs.values[4 + 6 * i])
            elem.raw_id = int(gs.values[5 + 6 * i]) if gs.values[5 + 6 * i] is not None else None
            output.elements.append(elem)
    return output


def _interpolate_traffic_sample(s1: TrafficSample, w1: float, s2: TrafficSample, w2: float) -> TrafficSample:
    output = TrafficSample()
    output.time = bi.time

    no_id = False
    for elem in s1.elements:
        if elem.id <= 0:
            no_id = True
            break
    for elem in s2.elements:
        if elem.id <= 0:
            no_id = True
            break
    if no_id:
        return s1 if w1 > w2 else s2

    for o1 in s1.elements:
        for o2 in s2.elements:
            if o1.id == o2.id and type(o1) == type(o2):
                elem = None
                if isinstance(o1, SpeedLimit) and isinstance(o2, SpeedLimit):
                    elem = SpeedLimit()
                    elem.id = o1.id
                    elem.raw_id = o1.raw_id if w1 > w2 else o2.raw_id
                    elem.limit_speed = o1.limit_speed if w1 > w2 else o2.limit_speed
                    if o1.pos_x is not None and o2.pos_x is not None:
                        elem.pos_x = o1.pos_x * w1 + o2.pos_x * w2
                    if o1.pos_y is not None and o2.pos_y is not None:
                        elem.pos_y = o1.pos_y * w1 + o2.pos_y * w2
                    if o1.pos_z is not None and o2.pos_z is not None:
                        elem.pos_z = o1.pos_z * w1 + o2.pos_z * w2
                elif isinstance(o1, TrafficSign) and isinstance(o2, TrafficSign):
                    elem = TrafficSign()
                    elem.id = o1.id
                    elem.raw_id = o1.raw_id if w1 > w2 else o2.raw_id
                    elem.sign_type = o1.sign_type if w1 > w2 else o2.sign_type
                    if o1.pos_x is not None and o2.pos_x is not None:
                        elem.pos_x = o1.pos_x * w1 + o2.pos_x * w2
                    if o1.pos_y is not None and o2.pos_y is not None:
                        elem.pos_y = o1.pos_y * w1 + o2.pos_y * w2
                    if o1.pos_z is not None and o2.pos_z is not None:
                        elem.pos_z = o1.pos_z * w1 + o2.pos_z * w2
                elif isinstance(o1, TrafficLight) and isinstance(o2, TrafficLight):
                    elem = TrafficLight()
                    elem.id = o1.id
                    elem.raw_id = o1.raw_id if w1 > w2 else o2.raw_id
                    elem.state = o1.state if w1 > w2 else o2.state
                    elem.arrows = o1.arrows if w1 > w2 else o2.arrows
                    if o1.pos_x is not None and o2.pos_x is not None:
                        elem.pos_x = o1.pos_x * w1 + o2.pos_x * w2
                    if o1.pos_y is not None and o2.pos_y is not None:
                        elem.pos_y = o1.pos_y * w1 + o2.pos_y * w2
                    if o1.pos_z is not None and o2.pos_z is not None:
                        elem.pos_z = o1.pos_z * w1 + o2.pos_z * w2
                    if o1.stop_pos_x is not None and o2.stop_pos_x is not None:
                        elem.stop_pos_x = o1.stop_pos_x * w1 + o2.stop_pos_x * w2
                    if o1.stop_pos_y is not None and o2.stop_pos_y is not None:
                        elem.stop_pos_y = o1.stop_pos_y * w1 + o2.stop_pos_y * w2
                elif isinstance(o1, StopLine) and isinstance(o2, StopLine):
                    elem = StopLine()
                    elem.id = o1.id
                    elem.raw_id = o1.raw_id if w1 > w2 else o2.raw_id
                    elem.distance = o1.distance * w1 + o2.distance * w2
                    if o1.heading is not None and o2.heading is not None:
                        elem.heading = o1.heading * w1 + o2.heading * w2
                elif isinstance(o1, SpeedBump) and isinstance(o2, SpeedBump):
                    elem = SpeedBump()
                    elem.id = o1.id
                    elem.raw_id = o1.raw_id if w1 > w2 else o2.raw_id
                    elem.distance = o1.distance * w1 + o2.distance * w2
                    if o1.heading is not None and o2.heading is not None:
                        elem.heading = o1.heading * w1 + o2.heading * w2
                elif isinstance(o1, ArrowMark) and isinstance(o2, ArrowMark):
                    elem = ArrowMark()
                    elem.id = o1.id
                    elem.raw_id = o1.raw_id if w1 > w2 else o2.raw_id
                    elem.arrows = o1.arrows if w1 > w2 else o2.arrows
                    if o1.pos_x is not None and o2.pos_x is not None:
                        elem.pos_x = o1.pos_x * w1 + o2.pos_x * w2
                    if o1.pos_y is not None and o2.pos_y is not None:
                        elem.pos_y = o1.pos_y * w1 + o2.pos_y * w2
                elif isinstance(o1, ProhibitionZone) and isinstance(o2, ProhibitionZone):
                    elem = ProhibitionZone()
                    elem.id = o1.id
                    elem.raw_id = o1.raw_id if w1 > w2 else o2.raw_id
                    elem.center_x = o1.center_x * w1 + o2.center_x * w2
                    elem.center_y = o1.center_y * w1 + o2.center_y * w2
                    elem.width = o1.width * w1 + o2.width * w2
                    elem.length = o1.length * w1 + o2.length * w2
                    if o1.heading is not None and o2.heading is not None:
                        elem.heading = o1.heading * w1 + o2.heading * w2
                elif isinstance(o1, ParkingSlot) and isinstance(o2, ParkingSlot):
                    elem = ParkingSlot()
                    elem.id = o1.id
                    elem.raw_id = o1.raw_id if w1 > w2 else o2.raw_id
                    elem.center_x = o1.center_x * w1 + o2.center_x * w2
                    elem.center_y = o1.center_y * w1 + o2.center_y * w2
                    elem.width = o1.width * w1 + o2.width * w2
                    elem.length = o1.length * w1 + o2.length * w2
                    if o1.heading is not None and o2.heading is not None:
                        elem.heading = o1.heading * w1 + o2.heading * w2
                    elem.occupancy = o1.occupancy if w1 > w2 else o2.occupancy
                elif isinstance(o1, ZebraCrossing) and isinstance(o2, ZebraCrossing):
                    elem = ZebraCrossing()
                    elem.id = o1.id
                    elem.raw_id = o1.raw_id if w1 > w2 else o2.raw_id
                    elem.near_distance = o1.near_distance * w1 + o2.near_distance * w2
                    if o1.far_distance is not None and o2.far_distance is not None:
                        elem.far_distance = o1.far_distance * w1 + o2.far_distance * w2
                    if o1.heading is not None and o2.heading is not None:
                        elem.heading = o1.heading * w1 + o2.heading * w2

                if elem is not None:
                    output.elements.append(elem)
                break

    return output


# Query TrafficSample for input / 获取TrafficSample，用于样本输入
def get_traffic_sample(channel: int) -> TrafficSample | None:
    s1 = None
    s2 = None
    w1 = 0.0
    w2 = 0.0
    protocol_id_v1 = 'traffic-sample-v1@' + str(channel)
    protocol_id_v2 = 'traffic-sample-v2@' + str(channel)
    if protocol_id_v2 in bi.input_samples:
        pair = bi.input_samples[protocol_id_v2]
        s1 = _conv_traffic_sample_v2(pair.sample1)
        w1 = pair.weight1
        s2 = _conv_traffic_sample_v2(pair.sample2)
        w2 = pair.weight2
    elif protocol_id_v1 in bi.input_samples:
        pair = bi.input_samples[protocol_id_v1]
        s1 = _conv_traffic_sample_v1(pair.sample1)
        w1 = pair.weight1
        s2 = _conv_traffic_sample_v1(pair.sample2)
        w2 = pair.weight2
    if s1 is not None and s2 is not None:
        return _interpolate_traffic_sample(s1, w1, s2, w2)
    return None
